{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7765UFHoyGx6"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "KsOkK8O69PyT"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RKQpW0JqQQmY"
      },
      "source": [
        "# Shape Constraints with Tensorflow Lattice\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r61fkA2i9Y3_"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/lattice/tutorials/shape_constraints\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/lattice/blob/master/docs/tutorials/shape_constraints.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/lattice/blob/master/docs/tutorials/shape_constraints.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/lattice/docs/tutorials/shape_constraints.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2plcL3iTVjsp"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial is an overview of the constraints and regularizers provided by the TensorFlow Lattice (TFL) library. Here we use TFL canned estimators on synthetic datasets, but note that everything in this tutorial can also be done with models constructed from TFL Keras layers.\n",
        "\n",
        "Before proceeding, make sure your runtime has all required packages installed (as imported in the code cells below)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x769lI12IZXB"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fbBVAR6UeRN5"
      },
      "source": [
        "Installing TF Lattice package:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bpXjJKpSd3j4"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "!pip install tensorflow-lattice"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jSVl9SHTeSGX"
      },
      "source": [
        "Importing required packages:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "both",
        "id": "iY6awAl058TV"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "from IPython.core.pylabtools import figsize\n",
        "import itertools\n",
        "import logging\n",
        "import matplotlib\n",
        "from matplotlib import pyplot as plt\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import sys\n",
        "import tensorflow_lattice as tfl\n",
        "logging.disable(sys.maxsize)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7TmBk_IGgJF0"
      },
      "source": [
        "Default values used in this guide:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kQHPyPsPUF92"
      },
      "outputs": [],
      "source": [
        "NUM_EPOCHS = 1000\n",
        "BATCH_SIZE = 64\n",
        "LEARNING_RATE=0.01"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FjR7D8Ag3z0d"
      },
      "source": [
        "## Training Dataset for Ranking Restaurants"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a1YetzbdFOij"
      },
      "source": [
        "Imagine a simplified scenario where we want to determine whether or not users will click on a restaurant search result. The task is to predict the clickthrough rate (CTR) given input features:\n",
        "- Average rating (`avg_rating`): a numeric feature with values in the range [1,5].\n",
        "- Number of reviews (`num_reviews`): a numeric feature with values capped at 200, which we use as a measure of trendiness.\n",
        "- Dollar rating (`dollar_rating`): a categorical feature with string values in the set {\"D\", \"DD\", \"DDD\", \"DDDD\"}.\n",
        "\n",
        "Here we create a synthetic dataset where the true CTR is given by the formula:\n",
        "$$\n",
        "CTR = 1 / (1 + exp\\{\\mbox{b(dollar_rating)}-\\mbox{avg_rating}\\times log(\\mbox{num_reviews}) /4 \\})\n",
        "$$\n",
        "where $b(\\cdot)$ translates each `dollar_rating` to a baseline value:\n",
        "$$\n",
        "\\mbox{D}\\to 3,\\ \\mbox{DD}\\to 2,\\ \\mbox{DDD}\\to 4,\\ \\mbox{DDDD}\\to 4.5. \n",
        "$$\n",
        "\n",
        "This formula reflects typical user patterns. e.g. given everything else fixed, users prefer restaurants with higher star ratings, and \"\\\\$\\\\$\" restaurants will receive more clicks than \"\\\\$\", followed by \"\\\\$\\\\$\\\\$\" and \"\\\\$\\\\$\\\\$\\\\$\". "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mKovnyv1jATw"
      },
      "outputs": [],
      "source": [
        "def click_through_rate(avg_ratings, num_reviews, dollar_ratings):\n",
        "  dollar_rating_baseline = {\"D\": 3, \"DD\": 2, \"DDD\": 4, \"DDDD\": 4.5}\n",
        "  return 1 / (1 + np.exp(\n",
        "      np.array([dollar_rating_baseline[d] for d in dollar_ratings]) -\n",
        "      avg_ratings * np.log1p(num_reviews) / 4))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BPlgRdt6jAbP"
      },
      "source": [
        "Let's take a look at the contour plots of this CTR function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KC5qX_XKmc7g"
      },
      "outputs": [],
      "source": [
        "def color_bar():\n",
        "  bar = matplotlib.cm.ScalarMappable(\n",
        "      norm=matplotlib.colors.Normalize(0, 1, True),\n",
        "      cmap=\"viridis\",\n",
        "  )\n",
        "  bar.set_array([0, 1])\n",
        "  return bar\n",
        "\n",
        "\n",
        "def plot_fns(fns, split_by_dollar=False, res=25):\n",
        "  \"\"\"Generates contour plots for a list of (name, fn) functions.\"\"\"\n",
        "  num_reviews, avg_ratings = np.meshgrid(\n",
        "      np.linspace(0, 200, num=res),\n",
        "      np.linspace(1, 5, num=res),\n",
        "  )\n",
        "  if split_by_dollar:\n",
        "    dollar_rating_splits = [\"D\", \"DD\", \"DDD\", \"DDDD\"]\n",
        "  else:\n",
        "    dollar_rating_splits = [None]\n",
        "  if len(fns) == 1:\n",
        "    fig, axes = plt.subplots(2, 2, sharey=True, tight_layout=False)\n",
        "  else:\n",
        "    fig, axes = plt.subplots(\n",
        "        len(dollar_rating_splits), len(fns), sharey=True, tight_layout=False)\n",
        "  axes = axes.flatten()\n",
        "  axes_index = 0\n",
        "  for dollar_rating_split in dollar_rating_splits:\n",
        "    for title, fn in fns:\n",
        "      if dollar_rating_split is not None:\n",
        "        dollar_ratings = np.repeat(dollar_rating_split, res**2)\n",
        "        values = fn(avg_ratings.flatten(), num_reviews.flatten(),\n",
        "                    dollar_ratings)\n",
        "        title = \"{}: dollar_rating={}\".format(title, dollar_rating_split)\n",
        "      else:\n",
        "        values = fn(avg_ratings.flatten(), num_reviews.flatten())\n",
        "      subplot = axes[axes_index]\n",
        "      axes_index += 1\n",
        "      subplot.contourf(\n",
        "          avg_ratings,\n",
        "          num_reviews,\n",
        "          np.reshape(values, (res, res)),\n",
        "          vmin=0,\n",
        "          vmax=1)\n",
        "      subplot.title.set_text(title)\n",
        "      subplot.set(xlabel=\"Average Rating\")\n",
        "      subplot.set(ylabel=\"Number of Reviews\")\n",
        "      subplot.set(xlim=(1, 5))\n",
        "\n",
        "  _ = fig.colorbar(color_bar(), cax=fig.add_axes([0.95, 0.2, 0.01, 0.6]))\n",
        "\n",
        "\n",
        "figsize(11, 11)\n",
        "plot_fns([(\"CTR\", click_through_rate)], split_by_dollar=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ol91olp3muNN"
      },
      "source": [
        "### Preparing Data\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H8BOshZS9xwn"
      },
      "source": [
        "We now need to create our synthetic datasets. We start by generating a simulated dataset of restaurants and their features."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MhqcOPdTT_wj"
      },
      "outputs": [],
      "source": [
        "def sample_restaurants(n):\n",
        "  avg_ratings = np.random.uniform(1.0, 5.0, n)\n",
        "  num_reviews = np.round(np.exp(np.random.uniform(0.0, np.log(200), n)))\n",
        "  dollar_ratings = np.random.choice([\"D\", \"DD\", \"DDD\", \"DDDD\"], n)\n",
        "  ctr_labels = click_through_rate(avg_ratings, num_reviews, dollar_ratings)\n",
        "  return avg_ratings, num_reviews, dollar_ratings, ctr_labels\n",
        "\n",
        "\n",
        "np.random.seed(42)\n",
        "avg_ratings, num_reviews, dollar_ratings, ctr_labels = sample_restaurants(2000)\n",
        "\n",
        "figsize(5, 5)\n",
        "fig, axs = plt.subplots(1, 1, sharey=False, tight_layout=False)\n",
        "for rating, marker in [(\"D\", \"o\"), (\"DD\", \"^\"), (\"DDD\", \"+\"), (\"DDDD\", \"x\")]:\n",
        "  plt.scatter(\n",
        "      x=avg_ratings[np.where(dollar_ratings == rating)],\n",
        "      y=num_reviews[np.where(dollar_ratings == rating)],\n",
        "      c=ctr_labels[np.where(dollar_ratings == rating)],\n",
        "      vmin=0,\n",
        "      vmax=1,\n",
        "      marker=marker,\n",
        "      label=rating)\n",
        "plt.xlabel(\"Average Rating\")\n",
        "plt.ylabel(\"Number of Reviews\")\n",
        "plt.legend()\n",
        "plt.xlim((1, 5))\n",
        "plt.title(\"Distribution of restaurants\")\n",
        "_ = fig.colorbar(color_bar(), cax=fig.add_axes([0.95, 0.2, 0.01, 0.6]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tRetsfLv_JSR"
      },
      "source": [
        "Let's produce the training, validation and testing datasets. When a restaurant is viewed in the search results, we can record user's engagement (click or no click) as a sample point. \n",
        "\n",
        "In practice, users often do not go through all search results. This means that users will likely only see restaurants already considered \"good\" by the current ranking model in use. As a result, \"good\" restaurants are more frequently impressed and over-represented in the training datasets. When using more features, the training dataset can have large gaps in \"bad\" parts of the feature space.\n",
        "\n",
        "When the model is used for ranking, it is often evaluated on all relevant results with a more uniform distribution that is not well-represented by the training dataset. A flexible and complicated model might fail in this case due to overfitting the over-represented data points and thus lack generalizability. We handle this issue by applying domain knowledge to add *shape constraints* that guide the model to make reasonable predictions when it cannot pick them up from the training dataset.\n",
        "\n",
        "In this example, the training dataset mostly consists of user interactions with good and popular restaurants. The testing dataset has a uniform distribution to simulate the evaluation setting discussed above. Note that such testing dataset will not be available in a real problem setting."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jS6WOtXQ8jwX"
      },
      "outputs": [],
      "source": [
        "def sample_dataset(n, testing_set):\n",
        "  (avg_ratings, num_reviews, dollar_ratings, ctr_labels) = sample_restaurants(n)\n",
        "  if testing_set:\n",
        "    # Testing has a more uniform distribution over all restaurants.\n",
        "    num_views = np.random.poisson(lam=3, size=n)\n",
        "  else:\n",
        "    # Training/validation datasets have more views on popular restaurants.\n",
        "    num_views = np.random.poisson(lam=ctr_labels * num_reviews / 50.0, size=n)\n",
        "\n",
        "  return pd.DataFrame({\n",
        "      \"avg_rating\": np.repeat(avg_ratings, num_views),\n",
        "      \"num_reviews\": np.repeat(num_reviews, num_views),\n",
        "      \"dollar_rating\": np.repeat(dollar_ratings, num_views),\n",
        "      \"clicked\": np.random.binomial(n=1, p=np.repeat(ctr_labels, num_views))\n",
        "  })\n",
        "\n",
        "\n",
        "# Generate datasets.\n",
        "np.random.seed(42)\n",
        "data_train = sample_dataset(500, testing_set=False)\n",
        "data_val = sample_dataset(500, testing_set=False)\n",
        "data_test = sample_dataset(500, testing_set=True)\n",
        "\n",
        "# Plotting dataset densities.\n",
        "figsize(12, 5)\n",
        "fig, axs = plt.subplots(1, 2, sharey=False, tight_layout=False)\n",
        "for ax, data, title in [(axs[0], data_train, \"training\"),\n",
        "                        (axs[1], data_test, \"testing\")]:\n",
        "  _, _, _, density = ax.hist2d(\n",
        "      x=data[\"avg_rating\"],\n",
        "      y=data[\"num_reviews\"],\n",
        "      bins=(np.linspace(1, 5, num=21), np.linspace(0, 200, num=21)),\n",
        "      density=True,\n",
        "      cmap=\"Blues\",\n",
        "  )\n",
        "  ax.set(xlim=(1, 5))\n",
        "  ax.set(ylim=(0, 200))\n",
        "  ax.set(xlabel=\"Average Rating\")\n",
        "  ax.set(ylabel=\"Number of Reviews\")\n",
        "  ax.title.set_text(\"Density of {} examples\".format(title))\n",
        "  _ = fig.colorbar(density, ax=ax)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4fVyLgpCT1nW"
      },
      "source": [
        "Defining input_fns used for training and evaluation:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DYzRTRR2GKoS"
      },
      "outputs": [],
      "source": [
        "train_input_fn = tf.compat.v1.estimator.inputs.pandas_input_fn(\n",
        "    x=data_train,\n",
        "    y=data_train[\"clicked\"],\n",
        "    batch_size=BATCH_SIZE,\n",
        "    num_epochs=NUM_EPOCHS,\n",
        "    shuffle=False,\n",
        ")\n",
        "\n",
        "# feature_analysis_input_fn is used for TF Lattice estimators.\n",
        "feature_analysis_input_fn = tf.compat.v1.estimator.inputs.pandas_input_fn(\n",
        "    x=data_train,\n",
        "    y=data_train[\"clicked\"],\n",
        "    batch_size=BATCH_SIZE,\n",
        "    num_epochs=1,\n",
        "    shuffle=False,\n",
        ")\n",
        "\n",
        "val_input_fn = tf.compat.v1.estimator.inputs.pandas_input_fn(\n",
        "    x=data_val,\n",
        "    y=data_val[\"clicked\"],\n",
        "    batch_size=BATCH_SIZE,\n",
        "    num_epochs=1,\n",
        "    shuffle=False,\n",
        ")\n",
        "\n",
        "test_input_fn = tf.compat.v1.estimator.inputs.pandas_input_fn(\n",
        "    x=data_test,\n",
        "    y=data_test[\"clicked\"],\n",
        "    batch_size=BATCH_SIZE,\n",
        "    num_epochs=1,\n",
        "    shuffle=False,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qoTrw3FZqvPK"
      },
      "source": [
        "## Fitting Gradient Boosted Trees"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZklNowexE3wB"
      },
      "source": [
        "Let's start off with only two features: `avg_rating` and `num_reviews`.\n",
        "\n",
        "We create a few auxillary functions for plotting and calculating validation and test metrics."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SX6rARJWURWl"
      },
      "outputs": [],
      "source": [
        "def analyze_two_d_estimator(estimator, name):\n",
        "  # Extract validation metrics.\n",
        "  metric = estimator.evaluate(input_fn=val_input_fn)\n",
        "  print(\"Validation AUC: {}\".format(metric[\"auc\"]))\n",
        "  metric = estimator.evaluate(input_fn=test_input_fn)\n",
        "  print(\"Testing AUC: {}\".format(metric[\"auc\"]))\n",
        "\n",
        "  def two_d_pred(avg_ratings, num_reviews):\n",
        "    results = estimator.predict(\n",
        "        tf.compat.v1.estimator.inputs.pandas_input_fn(\n",
        "            x=pd.DataFrame({\n",
        "                \"avg_rating\": avg_ratings,\n",
        "                \"num_reviews\": num_reviews,\n",
        "            }),\n",
        "            shuffle=False,\n",
        "        ))\n",
        "    return [x[\"logistic\"][0] for x in results]\n",
        "\n",
        "  def two_d_click_through_rate(avg_ratings, num_reviews):\n",
        "    return np.mean([\n",
        "        click_through_rate(avg_ratings, num_reviews,\n",
        "                           np.repeat(d, len(avg_ratings)))\n",
        "        for d in [\"D\", \"DD\", \"DDD\", \"DDDD\"]\n",
        "    ],\n",
        "                   axis=0)\n",
        "\n",
        "  figsize(11, 5)\n",
        "  plot_fns([(\"{} Estimated CTR\".format(name), two_d_pred),\n",
        "            (\"CTR\", two_d_click_through_rate)],\n",
        "           split_by_dollar=False)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JVef4f8yUUbs"
      },
      "source": [
        "We can fit TensorFlow gradient boosted decision trees on the dataset:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DnPYlRAo2mnQ"
      },
      "outputs": [],
      "source": [
        "feature_columns = [\n",
        "    tf.feature_column.numeric_column(\"num_reviews\"),\n",
        "    tf.feature_column.numeric_column(\"avg_rating\"),\n",
        "]\n",
        "gbt_estimator = tf.estimator.BoostedTreesClassifier(\n",
        "    feature_columns=feature_columns,\n",
        "    # Hyper-params optimized on validation set.\n",
        "    n_batches_per_layer=1,\n",
        "    max_depth=2,\n",
        "    n_trees=50,\n",
        "    learning_rate=0.05,\n",
        "    config=tf.estimator.RunConfig(tf_random_seed=42),\n",
        ")\n",
        "gbt_estimator.train(input_fn=train_input_fn)\n",
        "analyze_two_d_estimator(gbt_estimator, \"GBT\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nYZtd6YvsNdn"
      },
      "source": [
        "Even though the model has captured the general shape of the true CTR and has decent validation metrics, it has counter-intuitive behavior in several parts of the input space: the estimated CTR decreases as the average rating or number of reviews increase. This is due to a lack of sample points in areas not well-covered by the training dataset. The model simply has no way to deduce the correct behaviour solely from the data.\n",
        "\n",
        "To solve this issue, we enforce the shape constraint that the model must output values monotonically increasing with respect to both the average rating and the number of reviews. We will later see how to implement this in TFL.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Uf7WqGooFiEp"
      },
      "source": [
        "## Fitting a DNN"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_s2aT3x0E_tF"
      },
      "source": [
        "We can repeat the same steps with a DNN classifier. We can observe a similar pattern: not having enough sample points with small number of reviews results in nonsensical extrapolation. Note that even though the validation metric is better than the tree solution, the testing metric is much worse."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gFUeG6kLDNhO"
      },
      "outputs": [],
      "source": [
        "feature_columns = [\n",
        "    tf.feature_column.numeric_column(\"num_reviews\"),\n",
        "    tf.feature_column.numeric_column(\"avg_rating\"),\n",
        "]\n",
        "dnn_estimator = tf.estimator.DNNClassifier(\n",
        "    feature_columns=feature_columns,\n",
        "    # Hyper-params optimized on validation set.\n",
        "    hidden_units=[16, 8, 8],\n",
        "    optimizer=tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE),\n",
        "    config=tf.estimator.RunConfig(tf_random_seed=42),\n",
        ")\n",
        "dnn_estimator.train(input_fn=train_input_fn)\n",
        "analyze_two_d_estimator(dnn_estimator, \"DNN\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0Avkw-okw7JL"
      },
      "source": [
        "## Shape Constraints"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3ExyethCFBrP"
      },
      "source": [
        "TensorFlow Lattice (TFL) is focused on enforcing shape constraints to safeguard model behavior beyond the training data. These shape constraints are applied to TFL Keras layers. Their details can be found in [our JMLR paper](http://jmlr.org/papers/volume17/15-243/15-243.pdf). \n",
        "\n",
        "In this tutorial we use TF canned estimators to cover various shape constraints, but note that all these steps can be done with models created from TFL Keras layers.\n",
        "\n",
        "As with any other TensorFlow estimator, TFL canned estimators use [feature columns](https://www.tensorflow.org/api_docs/python/tf/feature_column) to define the input format and use a training input_fn to pass in the data.\n",
        "Using TFL canned estimators also requires:\n",
        "- a *model config*: defining the model architecture and per-feature shape constraints and regularizers.\n",
        "- a *feature analysis input_fn*: a TF input_fn passing data for TFL initialization.\n",
        "\n",
        "For a more thorough description, please refer to the canned estimators tutorial or the API docs."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "anyCM4sCpOSo"
      },
      "source": [
        "### Monotonicity\n",
        "We first address the monotonicity concerns by adding monotonicity shape constraints to both features.\n",
        "\n",
        "To instruct TFL to enforce shape constraints, we specify the constraints in the *feature configs*. The following code shows how we can require the output to be monotonically increasing with respect to both `num_reviews` and `avg_rating` by setting `monotonicity=\"increasing\"`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FCm1lOjmwur_"
      },
      "outputs": [],
      "source": [
        "feature_columns = [\n",
        "    tf.feature_column.numeric_column(\"num_reviews\"),\n",
        "    tf.feature_column.numeric_column(\"avg_rating\"),\n",
        "]\n",
        "model_config = tfl.configs.CalibratedLatticeConfig(\n",
        "    feature_configs=[\n",
        "        tfl.configs.FeatureConfig(\n",
        "            name=\"num_reviews\",\n",
        "            lattice_size=2,\n",
        "            monotonicity=\"increasing\",\n",
        "            pwl_calibration_num_keypoints=20,\n",
        "        ),\n",
        "        tfl.configs.FeatureConfig(\n",
        "            name=\"avg_rating\",\n",
        "            lattice_size=2,\n",
        "            monotonicity=\"increasing\",\n",
        "            pwl_calibration_num_keypoints=20,\n",
        "        )\n",
        "    ])\n",
        "tfl_estimator = tfl.estimators.CannedClassifier(\n",
        "    feature_columns=feature_columns,\n",
        "    model_config=model_config,\n",
        "    feature_analysis_input_fn=feature_analysis_input_fn,\n",
        "    optimizer=tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE),\n",
        "    config=tf.estimator.RunConfig(tf_random_seed=42),\n",
        ")\n",
        "tfl_estimator.train(input_fn=train_input_fn)\n",
        "analyze_two_d_estimator(tfl_estimator, \"TF Lattice\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ubNRBCWW5wQ9"
      },
      "source": [
        "Using a `CalibratedLatticeConfig` creates a canned classifier that first applies a *calibrator* to each input (a piece-wise linear function for numeric features) followed by a *lattice* layer to non-linearly fuse the calibrated features. We can use `tfl.visualization` to visualize the model. In particular, the following plot shows the two trained calibrators included in the canned classifier. \n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "both",
        "id": "C0py9Q6OBRBE"
      },
      "outputs": [],
      "source": [
        "def save_and_visualize_lattice(tfl_estimator):\n",
        "  saved_model_path = tfl_estimator.export_saved_model(\n",
        "      \"/tmp/TensorFlow_Lattice_101/\",\n",
        "      tf.estimator.export.build_parsing_serving_input_receiver_fn(\n",
        "          feature_spec=tf.feature_column.make_parse_example_spec(\n",
        "              feature_columns)))\n",
        "  model_graph = tfl.estimators.get_model_graph(saved_model_path)\n",
        "  figsize(8, 8)\n",
        "  tfl.visualization.draw_model_graph(model_graph)\n",
        "  return model_graph\n",
        "\n",
        "_ = save_and_visualize_lattice(tfl_estimator)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7vZ5fShXs504"
      },
      "source": [
        "With the constraints added, the estimated CTR will always increase as the average rating increases or the number of reviews increases. This is done by making sure that the calibrators and the lattice are monotonic."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RfniRZCHIvfK"
      },
      "source": [
        "### Diminishing Returns\n",
        "[Diminishing returns](https://en.wikipedia.org/wiki/Diminishing_returns) means that the marginal gain of increasing a certain feature value will decrease as we increase the value. In our case we expect that the `num_reviews` feature follows this pattern, so we can configure its calibrator accordingly. Notice that we can decompose diminishing returns into two sufficient conditions:\n",
        "\n",
        "- the calibrator is monotonicially increasing, and\n",
        "- the calibrator is concave.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XQrM9BskY-wx"
      },
      "outputs": [],
      "source": [
        "feature_columns = [\n",
        "    tf.feature_column.numeric_column(\"num_reviews\"),\n",
        "    tf.feature_column.numeric_column(\"avg_rating\"),\n",
        "]\n",
        "model_config = tfl.configs.CalibratedLatticeConfig(\n",
        "    feature_configs=[\n",
        "        tfl.configs.FeatureConfig(\n",
        "            name=\"num_reviews\",\n",
        "            lattice_size=2,\n",
        "            monotonicity=\"increasing\",\n",
        "            pwl_calibration_convexity=\"concave\",\n",
        "            pwl_calibration_num_keypoints=20,\n",
        "        ),\n",
        "        tfl.configs.FeatureConfig(\n",
        "            name=\"avg_rating\",\n",
        "            lattice_size=2,\n",
        "            monotonicity=\"increasing\",\n",
        "            pwl_calibration_num_keypoints=20,\n",
        "        )\n",
        "    ])\n",
        "tfl_estimator = tfl.estimators.CannedClassifier(\n",
        "    feature_columns=feature_columns,\n",
        "    model_config=model_config,\n",
        "    feature_analysis_input_fn=feature_analysis_input_fn,\n",
        "    optimizer=tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE),\n",
        "    config=tf.estimator.RunConfig(tf_random_seed=42),\n",
        ")\n",
        "tfl_estimator.train(input_fn=train_input_fn)\n",
        "analyze_two_d_estimator(tfl_estimator, \"TF Lattice\")\n",
        "_ = save_and_visualize_lattice(tfl_estimator)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LSmzHkPUo9u5"
      },
      "source": [
        "Notice how the testing metric improves by adding the concavity constraint. The prediction plot also better resembles the ground truth."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J6CP2Ovapiu3"
      },
      "source": [
        "### 2D Shape Constraint: Trust\n",
        "A 5-star rating for a restaurant with only one or two reviews is likely an unreliable rating (the restaurant might not actually be good), whereas a 4-star rating for a restaurant with hundreds of reviews is much more reliable (the restaurant is likely good in this case). We can see that the number of reviews of a restaurant affects how much trust we place in its average rating. \n",
        "\n",
        "We can exercise TFL trust constraints to inform the model that the larger (or smaller) value of one feature indicates more reliance or trust of another feature. This is done by setting `reflects_trust_in` configuration in the feature config."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OA14j0erm6TJ"
      },
      "outputs": [],
      "source": [
        "feature_columns = [\n",
        "    tf.feature_column.numeric_column(\"num_reviews\"),\n",
        "    tf.feature_column.numeric_column(\"avg_rating\"),\n",
        "]\n",
        "model_config = tfl.configs.CalibratedLatticeConfig(\n",
        "    feature_configs=[\n",
        "        tfl.configs.FeatureConfig(\n",
        "            name=\"num_reviews\",\n",
        "            lattice_size=2,\n",
        "            monotonicity=\"increasing\",\n",
        "            pwl_calibration_convexity=\"concave\",\n",
        "            pwl_calibration_num_keypoints=20,\n",
        "            # Larger num_reviews indicating more trust in avg_rating.\n",
        "            reflects_trust_in=[\n",
        "                tfl.configs.TrustConfig(\n",
        "                    feature_name=\"avg_rating\", trust_type=\"edgeworth\"),\n",
        "            ],\n",
        "        ),\n",
        "        tfl.configs.FeatureConfig(\n",
        "            name=\"avg_rating\",\n",
        "            lattice_size=2,\n",
        "            monotonicity=\"increasing\",\n",
        "            pwl_calibration_num_keypoints=20,\n",
        "        )\n",
        "    ])\n",
        "tfl_estimator = tfl.estimators.CannedClassifier(\n",
        "    feature_columns=feature_columns,\n",
        "    model_config=model_config,\n",
        "    feature_analysis_input_fn=feature_analysis_input_fn,\n",
        "    optimizer=tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE),\n",
        "    config=tf.estimator.RunConfig(tf_random_seed=42),\n",
        ")\n",
        "tfl_estimator.train(input_fn=train_input_fn)\n",
        "analyze_two_d_estimator(tfl_estimator, \"TF Lattice\")\n",
        "model_graph = save_and_visualize_lattice(tfl_estimator)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "puvP9X8XxyRV"
      },
      "source": [
        "The following plot presents the trained lattice function. Due to the trust constraint, we expect that larger values of calibrated `num_reviews` would force higher slope with respect to calibrated `avg_rating`, resulting in a more significant move in the lattice output."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "both",
        "id": "RounEQebxxnA"
      },
      "outputs": [],
      "source": [
        "lat_mesh_n = 12\n",
        "lat_mesh_x, lat_mesh_y = tfl.test_utils.two_dim_mesh_grid(\n",
        "    lat_mesh_n**2, 0, 0, 1, 1)\n",
        "lat_mesh_fn = tfl.test_utils.get_hypercube_interpolation_fn(\n",
        "    model_graph.output_node.weights.flatten())\n",
        "lat_mesh_z = [\n",
        "    lat_mesh_fn([lat_mesh_x.flatten()[i],\n",
        "                 lat_mesh_y.flatten()[i]]) for i in range(lat_mesh_n**2)\n",
        "]\n",
        "trust_plt = tfl.visualization.plot_outputs(\n",
        "    (lat_mesh_x, lat_mesh_y),\n",
        "    {\"Lattice Lookup\": lat_mesh_z},\n",
        "    figsize=(6, 6),\n",
        ")\n",
        "trust_plt.title(\"Trust\")\n",
        "trust_plt.xlabel(\"Calibrated avg_rating\")\n",
        "trust_plt.ylabel(\"Calibrated num_reviews\")\n",
        "trust_plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SKe3UHX6pUjw"
      },
      "source": [
        "### Smoothing Calibrators\n",
        "Let's now take a look at the calibrator of `avg_rating`. Though it is monotonically increasing, the changes in its slopes are abrupt and hard to interpret. That suggests we might want to consider smoothing this calibrator using a regularizer setup in the `regularizer_configs`.\n",
        "\n",
        "Here we apply a `wrinkle` regularizer to reduce changes in the curvature. You can also use the `laplacian` regularizer to flatten the calibrator and the `hessian` regularizer to make it more linear.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qxFHH3hSpWfq"
      },
      "outputs": [],
      "source": [
        "feature_columns = [\n",
        "    tf.feature_column.numeric_column(\"num_reviews\"),\n",
        "    tf.feature_column.numeric_column(\"avg_rating\"),\n",
        "]\n",
        "model_config = tfl.configs.CalibratedLatticeConfig(\n",
        "    feature_configs=[\n",
        "        tfl.configs.FeatureConfig(\n",
        "            name=\"num_reviews\",\n",
        "            lattice_size=2,\n",
        "            monotonicity=\"increasing\",\n",
        "            pwl_calibration_convexity=\"concave\",\n",
        "            pwl_calibration_num_keypoints=20,\n",
        "            regularizer_configs=[\n",
        "                tfl.configs.RegularizerConfig(name=\"calib_wrinkle\", l2=1.0),\n",
        "            ],\n",
        "            reflects_trust_in=[\n",
        "                tfl.configs.TrustConfig(\n",
        "                    feature_name=\"avg_rating\", trust_type=\"edgeworth\"),\n",
        "            ],\n",
        "        ),\n",
        "        tfl.configs.FeatureConfig(\n",
        "            name=\"avg_rating\",\n",
        "            lattice_size=2,\n",
        "            monotonicity=\"increasing\",\n",
        "            pwl_calibration_num_keypoints=20,\n",
        "            regularizer_configs=[\n",
        "                tfl.configs.RegularizerConfig(name=\"calib_wrinkle\", l2=1.0),\n",
        "            ],\n",
        "        )\n",
        "    ])\n",
        "tfl_estimator = tfl.estimators.CannedClassifier(\n",
        "    feature_columns=feature_columns,\n",
        "    model_config=model_config,\n",
        "    feature_analysis_input_fn=feature_analysis_input_fn,\n",
        "    optimizer=tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE),\n",
        "    config=tf.estimator.RunConfig(tf_random_seed=42),\n",
        ")\n",
        "tfl_estimator.train(input_fn=train_input_fn)\n",
        "analyze_two_d_estimator(tfl_estimator, \"TF Lattice\")\n",
        "_ = save_and_visualize_lattice(tfl_estimator)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HHpp4goLvuPi"
      },
      "source": [
        "The calibrators are now smooth, and the overall estimated CTR better matches the ground truth. This is reflected both in the testing metric and in the contour plots."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pSUd6aFlpYz4"
      },
      "source": [
        "### Partial Monotonicity for Categorical Calibration\n",
        "So far we have been using only two of the numeric features in the model. Here we will add a third feature using a categorical calibration layer. Again we start by setting up helper functions for plotting and metric calculation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5tLDKwTmjrLw"
      },
      "outputs": [],
      "source": [
        "def analyze_three_d_estimator(estimator, name):\n",
        "  # Extract validation metrics.\n",
        "  metric = estimator.evaluate(input_fn=val_input_fn)\n",
        "  print(\"Validation AUC: {}\".format(metric[\"auc\"]))\n",
        "  metric = estimator.evaluate(input_fn=test_input_fn)\n",
        "  print(\"Testing AUC: {}\".format(metric[\"auc\"]))\n",
        "\n",
        "  def three_d_pred(avg_ratings, num_reviews, dollar_rating):\n",
        "    results = estimator.predict(\n",
        "        tf.compat.v1.estimator.inputs.pandas_input_fn(\n",
        "            x=pd.DataFrame({\n",
        "                \"avg_rating\": avg_ratings,\n",
        "                \"num_reviews\": num_reviews,\n",
        "                \"dollar_rating\": dollar_rating,\n",
        "            }),\n",
        "            shuffle=False,\n",
        "        ))\n",
        "    return [x[\"logistic\"][0] for x in results]\n",
        "\n",
        "  figsize(11, 22)\n",
        "  plot_fns([(\"{} Estimated CTR\".format(name), three_d_pred),\n",
        "            (\"CTR\", click_through_rate)],\n",
        "           split_by_dollar=True)\n",
        "  "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CnPiqf4rq6kJ"
      },
      "source": [
        "To involve the third feature, `dollar_rating`, we should recall that categorical features require a slightly different treatment in TFL, both as a feature column and as a feature config. Here we enforce the partial monotonicity constraint that outputs for \"DD\" restaurants should be larger than \"D\" restaurants when all other inputs are fixed. This is done using the `monotonicity` setting in the feature config."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m-w7iGEEpgGt"
      },
      "outputs": [],
      "source": [
        "feature_columns = [\n",
        "    tf.feature_column.numeric_column(\"num_reviews\"),\n",
        "    tf.feature_column.numeric_column(\"avg_rating\"),\n",
        "    tf.feature_column.categorical_column_with_vocabulary_list(\n",
        "        \"dollar_rating\",\n",
        "        vocabulary_list=[\"D\", \"DD\", \"DDD\", \"DDDD\"],\n",
        "        dtype=tf.string,\n",
        "        default_value=0),\n",
        "]\n",
        "model_config = tfl.configs.CalibratedLatticeConfig(\n",
        "    feature_configs=[\n",
        "        tfl.configs.FeatureConfig(\n",
        "            name=\"num_reviews\",\n",
        "            lattice_size=2,\n",
        "            monotonicity=\"increasing\",\n",
        "            pwl_calibration_convexity=\"concave\",\n",
        "            pwl_calibration_num_keypoints=20,\n",
        "            regularizer_configs=[\n",
        "                tfl.configs.RegularizerConfig(name=\"calib_wrinkle\", l2=1.0),\n",
        "            ],\n",
        "            reflects_trust_in=[\n",
        "                tfl.configs.TrustConfig(\n",
        "                    feature_name=\"avg_rating\", trust_type=\"edgeworth\"),\n",
        "            ],\n",
        "        ),\n",
        "        tfl.configs.FeatureConfig(\n",
        "            name=\"avg_rating\",\n",
        "            lattice_size=2,\n",
        "            monotonicity=\"increasing\",\n",
        "            pwl_calibration_num_keypoints=20,\n",
        "            regularizer_configs=[\n",
        "                tfl.configs.RegularizerConfig(name=\"calib_wrinkle\", l2=1.0),\n",
        "            ],\n",
        "        ),\n",
        "        tfl.configs.FeatureConfig(\n",
        "            name=\"dollar_rating\",\n",
        "            lattice_size=2,\n",
        "            pwl_calibration_num_keypoints=4,\n",
        "            # Here we only specify one monotonicity:\n",
        "            # `D` resturants has smaller value than `DD` restaurants\n",
        "            monotonicity=[(\"D\", \"DD\")],\n",
        "        ),\n",
        "    ])\n",
        "tfl_estimator = tfl.estimators.CannedClassifier(\n",
        "    feature_columns=feature_columns,\n",
        "    model_config=model_config,\n",
        "    feature_analysis_input_fn=feature_analysis_input_fn,\n",
        "    optimizer=tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE),\n",
        "    config=tf.estimator.RunConfig(tf_random_seed=42),\n",
        ")\n",
        "tfl_estimator.train(input_fn=train_input_fn)\n",
        "analyze_three_d_estimator(tfl_estimator, \"TF Lattice\")\n",
        "_ = save_and_visualize_lattice(tfl_estimator)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gdIzhYL79_Pp"
      },
      "source": [
        "This categorical calibrator shows the preference of the model output: DD > D > DDD > DDDD, which is consistent with our setup. Notice there is also a column for missing values. Though there is no missing feature in our training and testing data, the model provides us with an imputation for the missing value should it happen during downstream model serving.\n",
        "\n",
        "Here we also plot the predicted CTR of this model conditioned on `dollar_rating`. Notice that all the constraints we required are fulfilled in each of the slices."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rh0H2b6l_rwZ"
      },
      "source": [
        "### Output Calibration"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KPb2ri4e7HXF"
      },
      "source": [
        "For all the TFL models we have trained so far, the lattice layer (indicated as \"Lattice\" in the model graph) directly outputs the model prediction. Sometimes we are not sure whether the lattice output should be rescaled to emit model outputs:\n",
        "- the features are $log$ counts while the labels are counts.\n",
        "- the lattice is configured to have very few vertices but the label distribution is relatively complicated.\n",
        "\n",
        "In those cases we can add another calibrator between the lattice output and the model output to increase model flexibility. Here let's add a calibrator layer with 5 keypoints to the model we just built. We also add a regularizer for the output calibrator to keep the function smooth.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "k5Sg_gUj_0i4"
      },
      "outputs": [],
      "source": [
        "feature_columns = [\n",
        "    tf.feature_column.numeric_column(\"num_reviews\"),\n",
        "    tf.feature_column.numeric_column(\"avg_rating\"),\n",
        "    tf.feature_column.categorical_column_with_vocabulary_list(\n",
        "        \"dollar_rating\",\n",
        "        vocabulary_list=[\"D\", \"DD\", \"DDD\", \"DDDD\"],\n",
        "        dtype=tf.string,\n",
        "        default_value=0),\n",
        "]\n",
        "model_config = tfl.configs.CalibratedLatticeConfig(\n",
        "    output_calibration=True,\n",
        "    output_calibration_num_keypoints=5,\n",
        "    regularizer_configs=[\n",
        "        tfl.configs.RegularizerConfig(name=\"output_calib_wrinkle\", l2=0.1),\n",
        "    ],\n",
        "    feature_configs=[\n",
        "    tfl.configs.FeatureConfig(\n",
        "        name=\"num_reviews\",\n",
        "        lattice_size=2,\n",
        "        monotonicity=\"increasing\",\n",
        "        pwl_calibration_convexity=\"concave\",\n",
        "        pwl_calibration_num_keypoints=20,\n",
        "        regularizer_configs=[\n",
        "            tfl.configs.RegularizerConfig(name=\"calib_wrinkle\", l2=1.0),\n",
        "        ],\n",
        "        reflects_trust_in=[\n",
        "            tfl.configs.TrustConfig(\n",
        "                feature_name=\"avg_rating\", trust_type=\"edgeworth\"),\n",
        "        ],\n",
        "    ),\n",
        "    tfl.configs.FeatureConfig(\n",
        "        name=\"avg_rating\",\n",
        "        lattice_size=2,\n",
        "        monotonicity=\"increasing\",\n",
        "        pwl_calibration_num_keypoints=20,\n",
        "        regularizer_configs=[\n",
        "            tfl.configs.RegularizerConfig(name=\"calib_wrinkle\", l2=1.0),\n",
        "        ],\n",
        "    ),\n",
        "    tfl.configs.FeatureConfig(\n",
        "        name=\"dollar_rating\",\n",
        "        lattice_size=2,\n",
        "        pwl_calibration_num_keypoints=4,\n",
        "        # Here we only specify one monotonicity:\n",
        "        # `D` resturants has smaller value than `DD` restaurants\n",
        "        monotonicity=[(\"D\", \"DD\")],\n",
        "    ),\n",
        "])\n",
        "tfl_estimator = tfl.estimators.CannedClassifier(\n",
        "    feature_columns=feature_columns,\n",
        "    model_config=model_config,\n",
        "    feature_analysis_input_fn=feature_analysis_input_fn,\n",
        "    optimizer=tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE),\n",
        "    config=tf.estimator.RunConfig(tf_random_seed=42),\n",
        ")\n",
        "tfl_estimator.train(input_fn=train_input_fn)\n",
        "analyze_three_d_estimator(tfl_estimator, \"TF Lattice\")\n",
        "_ = save_and_visualize_lattice(tfl_estimator)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TLOGDrYY0hH7"
      },
      "source": [
        "The final testing metric and plots show how using common-sense constraints can help the model avoid unexpected behaviour and extrapolate better to the entire input space."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "shape_constraints.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
